"""
多模态MCTS推理模块

该模块扩展了原有的MCTS推理系统以支持多模态输入，主要功能包括：
- 复用现有的MCTS框架和推理逻辑
- 扩展Generator类以支持图像输入
- 保持原有的prompt系统和推理动作（A1-A5）
- 优化并行调用和图像处理
- 最小化代码修改，避免冗余

核心设计原则：
- 继承和扩展现有的Generator类
- 复用所有现有的prompt模板和推理逻辑
- 只修改模型调用部分以支持多模态输入
- 保持与原有系统的完全兼容性
- 优化API并行调用性能
"""

import sys
sys.path.append(".")

from run_src.MCTS_for_reasoning import Generator, search_for_answers
from models.multimodal_model import MultimodalModelInterface
from common.config_loader import config_to_args
from typing import Dict, Any, Optional, List
import copy
import logging
import json

# 设置日志
logger = logging.getLogger(__name__)


class MultimodalGenerator(Generator):
    """
    支持多模态的推理生成器
    
    该类继承自原有的Generator类，扩展其功能以支持多模态输入。
    保持所有原有的推理逻辑、prompt模板和推理动作不变，只修改模型调用部分。
    
    主要特性：
    - 完全继承原有Generator的所有功能
    - 复用所有现有的prompt模板和推理逻辑
    - 支持图像和文本的联合推理
    - 保持与原有系统的完全兼容性
    - 优化并行调用和图像处理
    
    Attributes:
        model_interface (MultimodalModelInterface): 多模态模型接口
        current_image_path (Optional[str]): 当前处理的图像路径
    """
    
    def __init__(self, config, model_interface: MultimodalModelInterface):
        """
        初始化多模态推理生成器
        
        Args:
            config: 配置对象，包含各种配置选项
            model_interface (MultimodalModelInterface): 多模态模型接口
        """
        # 首先设置模型接口和基本属性
        self.model_interface = model_interface
        self.current_image_path: Optional[str] = None
        
        # 转换配置为args格式以兼容原有代码
        args = config_to_args(config)
        
        # 复制原有Generator的配置参数
        self.num_subquestions = args.num_subquestions
        self.num_a1_steps = args.num_a1_steps
        self.num_votes = args.num_votes
        self.mcts_num_last_votes = args.mcts_num_last_votes
        self.enable_potential_score = args.enable_potential_score
        self.max_tokens = args.max_tokens
        
        # 初始化评估器
        from eval_src.Evaluator import GSM8KEvaluator
        self.evaluator = GSM8KEvaluator()
        
        # 加载提示模板（复用原有逻辑）
        self._load_prompts_from_original(args)
        
        # 最后创建IO兼容层（确保model_interface已经设置）
        self.io = self._create_io_compatibility_layer()
        
        logger.info("多模态推理生成器初始化完成")
    
    def _create_io_compatibility_layer(self):
        """
        创建IO兼容层
        
        将多模态模型接口包装成原有Generator期望的IO接口格式。
        这样可以保持原有的推理逻辑不变，只改变底层的模型调用。
        
        Returns:
            object: 兼容的IO接口对象
        """
        # 调试信息
        logger.debug(f"创建IO兼容层 - self.model_interface: {self.model_interface}")
        logger.debug(f"创建IO兼容层 - self.model_interface type: {type(self.model_interface)}")
        
        class IOCompatibilityLayer:
            def __init__(self, model_interface, multimodal_generator):
                self.model_interface = model_interface
                self.multimodal_generator = multimodal_generator
                self.call_counter = 0
                self.token_counter = 0
                
            def generate(self, model_input, max_tokens: int, num_return: int = 1, stop_tokens=None, **kwargs) -> List[str]:
                """
                生成文本响应，兼容原有的IO接口
                
                Args:
                    model_input (str): 输入文本
                    max_tokens (int): 最大生成token数
                    num_return (int): 返回结果数量
                    stop_tokens: 停止token列表
                    **kwargs: 其他参数
                    
                Returns:
                    List[str]: 生成的响应列表
                """
                self.call_counter += 1
                
                # 获取当前图像路径
                image_path = self.multimodal_generator.current_image_path
                
                # 显示图像使用信息
                if hasattr(self.model_interface, 'show_image_info') and self.model_interface.show_image_info:
                    if image_path:
                        logger.info(f"🔄 第{self.call_counter}次调用 - 包含图像: {image_path}")
                    else:
                        logger.info(f"🔄 第{self.call_counter}次调用 - 纯文本推理")
                
                # 准备消息
                messages = [{'content': model_input}]
                image_paths = [image_path] if image_path else None
                
                # 调用多模态模型接口
                try:
                    results = self.model_interface.generate(
                        messages=messages,
                        image_paths=image_paths,
                        num_return=num_return
                    )
                    
                    # 确保返回正确数量的结果
                    while len(results) < num_return:
                        results.append("")
                    
                    return results[:num_return]
                    
                except Exception as e:
                    logger.error(f"多模态生成失败: {e}")
                    return [""] * num_return
            
            def get_token_count(self) -> int:
                """获取token计数"""
                return self.token_counter
            
            def get_call_count(self) -> int:
                """获取调用计数"""
                return self.call_counter
        
        return IOCompatibilityLayer(self.model_interface, self)
    
    def _load_prompts_from_original(self, args):
        """
        从原有Generator加载提示模板
        
        Args:
            args: 命令行参数对象
        """
        try:
            # 直接复制原有Generator的初始化逻辑
            # 加载问题索引配置
            with open(args.decompose_template_path, "r") as f:
                decompose_template = json.load(f)
                self.question_index = decompose_template["index"]

            # 加载基础提示模板
            from common.utils import read_txt, read_json
            self.decompose_prompt = read_txt(args.decompose_prompt_path)
            self.fewshot_cot_prompt = read_txt(args.fewshot_cot_prompt_path)
            self.fewshot_cot_config = read_json(args.fewshot_cot_config_path)

            # 加载一步思考相关提示（A1动作）
            if not args.disable_a1:
                self.fewshot_ost_prompt = read_txt(args.fewshot_ost_prompt_path)
                self.fewshot_ost_config = read_json(args.fewshot_ost_config_path)

            # 加载问题重述相关提示（A5动作）
            if not args.disable_a5:
                import os
                self.rephrasing_prompt_template = read_txt(args.rephrasing_prompt_template_path)
                self.decompose_prompt_rephrased = read_txt(args.decompose_prompt_rephrased_path)
                self.fewshot_cot_prompt_rephrased = read_txt(args.fewshot_cot_prompt_rephrased_path)
                # 检查文件是否存在，如果不存在则使用标准版本
                if hasattr(args, 'fewshot_ost_prompt_rephrased_path') and os.path.exists(args.fewshot_ost_prompt_rephrased_path):
                    self.fewshot_ost_prompt_rephrased = read_txt(args.fewshot_ost_prompt_rephrased_path)
                elif hasattr(args, 'fewshot_ost_prompt_path') and os.path.exists(args.fewshot_ost_prompt_path):
                    self.fewshot_ost_prompt_rephrased = read_txt(args.fewshot_ost_prompt_path)
            
            logger.info("提示模板加载完成")
            
        except Exception as e:
            logger.error(f"提示模板加载失败: {e}")
            # 设置默认值
            self.fewshot_cot_prompt = ""
            self.decompose_prompt = ""
            self.fewshot_cot_config = {}
            self.question_index = 0
    
    def set_current_image(self, image_path: Optional[str]):
        """
        设置当前处理的图像路径
        
        Args:
            image_path (Optional[str]): 图像路径
        """
        self.current_image_path = image_path
        if hasattr(self.model_interface, 'log_image_usage') and self.model_interface.log_image_usage:
            if image_path:
                logger.info(f"📸 设置当前图像: {image_path}")
            else:
                logger.info("📝 清除当前图像，切换到纯文本模式")


def multimodal_search_for_answers(
    config,
    question: str,
    image_path: Optional[str] = None,
    **kwargs
) -> Dict[str, Any]:
    """
    多模态MCTS搜索
    
    该函数是原有search_for_answers函数的多模态版本，支持图像和文本的联合推理。
    
    Args:
        config: 配置对象
        question (str): 问题文本
        image_path (Optional[str]): 图像路径
        **kwargs: 其他参数
        
    Returns:
        Dict[str, Any]: 搜索结果
    """
    logger.info("开始多模态MCTS搜索")
    logger.info(f"问题: {question[:100]}{'...' if len(question) > 100 else ''}")
    if image_path:
        logger.info(f"图像: {image_path}")
    else:
        logger.info("纯文本模式")
    
    # 初始化模型接口
    model_interface = MultimodalModelInterface(config)
    
    # 创建多模态生成器
    generator = MultimodalGenerator(config, model_interface)
    
    # 设置当前图像
    generator.set_current_image(image_path)
    
    # 转换配置为args格式
    args = config_to_args(config)
    
    # 调用原有的search_for_answers函数
    try:
        # 提取必要的参数
        question_id = kwargs.get('question_id', 0)
        gt_answer = kwargs.get('gt_answer', '')
        
        # 调用原有的search_for_answers函数
        model_solutions, stopping_id, model_all_solutions = search_for_answers(
            args=args,
            user_question=question,
            question_id=question_id,
            gt_answer=gt_answer,
            generator=generator
        )
        
        # 构建返回结果
        result = {
            'question': question,
            'question_id': question_id,
            'gt_answer': gt_answer,
            'model_solutions': model_solutions,
            'stopping_id': stopping_id,
            'model_all_solutions': model_all_solutions,
            'best_answer': model_solutions[-1] if model_solutions else None
        }
        
        # 添加多模态相关信息
        result.update({
            'image_path': image_path,
            'has_image': image_path is not None,
            'model_stats': model_interface.get_stats()
        })
        
        logger.info("多模态MCTS搜索完成")
        return result
        
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        logger.error(f"多模态MCTS搜索失败: {str(e)}")
        logger.error(f"详细错误信息: {error_details}")
        
        # 记录更多上下文信息
        logger.error(f"问题: {question[:200]}...")
        logger.error(f"图像路径: {image_path}")
        logger.error(f"问题ID: {kwargs.get('question_id', 'N/A')}")
        
        return {
            'question': question,
            'image_path': image_path,
            'has_image': image_path is not None,
            'error': str(e),
            'error_details': error_details,
            'best_answer': None,
            'search_tree': None
        }


def batch_multimodal_search(
    config,
    questions_and_images: List[Dict[str, Any]],
    **kwargs
) -> List[Dict[str, Any]]:
    """
    批量多模态MCTS搜索
    
    Args:
        config: 配置对象
        questions_and_images (List[Dict[str, Any]]): 问题和图像列表
        **kwargs: 其他参数
        
    Returns:
        List[Dict[str, Any]]: 搜索结果列表
    """
    results = []
    
    for i, item in enumerate(questions_and_images):
        logger.info(f"处理第 {i+1}/{len(questions_and_images)} 个样本")
        
        question = item.get('question', '')
        image_path = item.get('image_path')
        
        result = multimodal_search_for_answers(
            config=config,
            question=question,
            image_path=image_path,
            **kwargs
        )
        
        # 添加样本信息
        result.update({
            'sample_id': i,
            'original_data': item
        })
        
        results.append(result)
    
    return results 